"""
实时空调系统负荷预测模型
基于Transformer的多特征时间序列预测
支持实时更新和今日剩余时间预测
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Optional, Tuple, Dict
import datetime
from dataclasses import dataclass

# 导入基础Transformer组件
from labml_nn.transformers.models import Encoder, TransformerLayer, EmbeddingsWithLearnedPositionalEncoding
from labml_nn.transformers.mha import MultiHeadAttention
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.utils import subsequent_mask


@dataclass
class HVACForecastConfig:
    """空调负荷预测配置"""
    # 模型参数
    d_model: int = 128          # 模型维度
    n_heads: int = 8            # 注意力头数
    n_layers: int = 6           # Transformer层数
    d_ff: int = 512            # 前馈网络维度
    dropout: float = 0.1        # Dropout率
    
    # 时间参数
    seq_len: int = 168          # 输入序列长度（一周的小时数）
    max_pred_len: int = 24      # 最大预测长度（今日剩余小时数）
    update_interval: int = 1    # 更新间隔（小时）
    
    # 特征参数
    n_numerical_features: int = 6    # 数值特征数量
    n_categorical_features: int = 4   # 分类特征数量
    
    # 设备配置
    device: str = "cuda" if torch.cuda.is_available() else "cpu"


class TimeFeatureEmbedding(nn.Module):
    """时间特征嵌入层"""
    
    def __init__(self, d_model: int):
        super().__init__()
        self.hour_embed = nn.Embedding(24, d_model // 4)      # 小时嵌入
        self.day_embed = nn.Embedding(7, d_model // 4)        # 星期嵌入
        self.month_embed = nn.Embedding(12, d_model // 4)     # 月份嵌入
        self.season_embed = nn.Embedding(4, d_model // 4)     # 季节嵌入
        
    def forward(self, time_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            time_features: [batch_size, seq_len, 4] (hour, day, month, season)
        Returns:
            time_embed: [batch_size, seq_len, d_model]
        """
        hour = self.hour_embed(time_features[:, :, 0])
        day = self.day_embed(time_features[:, :, 1])
        month = self.month_embed(time_features[:, :, 2])
        season = self.season_embed(time_features[:, :, 3])
        
        return torch.cat([hour, day, month, season], dim=-1)


class MultiFeatureEmbedding(nn.Module):
    """多特征融合嵌入层"""
    
    def __init__(self, config: HVACForecastConfig):
        super().__init__()
        self.config = config
        
        # 负荷值投影
        self.load_projection = nn.Linear(1, config.d_model // 2)
        
        # 数值特征投影（温度、湿度、风速等）
        self.numerical_projection = nn.Linear(
            config.n_numerical_features, config.d_model // 4
        )
        
        # 时间特征嵌入
        self.time_embedding = TimeFeatureEmbedding(config.d_model // 4)
        
        # 特征融合层
        self.feature_fusion = nn.Linear(config.d_model, config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, load_data: torch.Tensor, 
                numerical_features: torch.Tensor,
                time_features: torch.Tensor) -> torch.Tensor:
        """
        Args:
            load_data: [batch_size, seq_len, 1] 历史负荷数据
            numerical_features: [batch_size, seq_len, n_numerical_features] 数值特征
            time_features: [batch_size, seq_len, 4] 时间特征
        Returns:
            embedded_features: [batch_size, seq_len, d_model]
        """
        # 各特征投影
        load_embed = self.load_projection(load_data)
        numerical_embed = self.numerical_projection(numerical_features)
        time_embed = self.time_embedding(time_features)
        
        # 特征拼接
        combined_features = torch.cat([load_embed, numerical_embed, time_embed], dim=-1)
        
        # 特征融合
        fused_features = self.feature_fusion(combined_features)
        
        return self.dropout(fused_features)


class RealTimeHVACTransformer(nn.Module):
    """实时空调负荷预测Transformer模型"""
    
    def __init__(self, config: HVACForecastConfig):
        super().__init__()
        self.config = config
        
        # 多特征嵌入层
        self.feature_embedding = MultiFeatureEmbedding(config)
        
        # 位置编码
        self.positional_encoding = EmbeddingsWithLearnedPositionalEncoding(
            config.d_model, 1, config.seq_len + config.max_pred_len
        )
        
        # Transformer编码器
        encoder_layer = TransformerLayer(
            d_model=config.d_model,
            self_attn=MultiHeadAttention(config.n_heads, config.d_model, config.dropout),
            src_attn=None,
            feed_forward=FeedForward(config.d_model, config.d_ff, config.dropout),
            dropout_prob=config.dropout
        )
        self.encoder = Encoder(encoder_layer, config.n_layers)
        
        # 输出投影层
        self.output_projection = nn.Sequential(
            nn.Linear(config.d_model, config.d_model // 2),
            nn.ReLU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.d_model // 2, 1)
        )
        
        # 缓存掩码以提高推理效率
        self.register_buffer('cached_mask', None)
        
    def create_mask(self, seq_len: int) -> torch.Tensor:
        """创建因果掩码"""
        if self.cached_mask is None or self.cached_mask.size(0) < seq_len:
            self.cached_mask = subsequent_mask(seq_len).to(self.config.device)
        return self.cached_mask[:seq_len, :seq_len]
    
    def forward(self, load_data: torch.Tensor,
                numerical_features: torch.Tensor,
                time_features: torch.Tensor,
                pred_len: int = None) -> torch.Tensor:
        """
        前向传播
        
        Args:
            load_data: [batch_size, seq_len, 1] 历史负荷数据
            numerical_features: [batch_size, seq_len, n_numerical_features] 数值特征
            time_features: [batch_size, seq_len, 4] 时间特征
            pred_len: 预测长度（默认为今日剩余小时数）
            
        Returns:
            predictions: [batch_size, pred_len, 1] 预测结果
        """
        batch_size, seq_len = load_data.shape[:2]
        
        if pred_len is None:
            # 计算今日剩余小时数
            current_hour = datetime.datetime.now().hour
            pred_len = 24 - current_hour
        
        # 特征嵌入
        embedded_features = self.feature_embedding(
            load_data, numerical_features, time_features
        )
        
        # 添加位置编码（仅用于序列位置信息）
        dummy_tokens = torch.zeros(batch_size, seq_len, dtype=torch.long, device=load_data.device)
        pos_encoding = self.positional_encoding(dummy_tokens)
        embedded_features = embedded_features + pos_encoding
        
        # 创建掩码
        mask = self.create_mask(seq_len)
        
        # Transformer编码
        encoded_features = self.encoder(embedded_features, mask)
        
        # 自回归预测
        predictions = []
        current_input = encoded_features
        
        for step in range(pred_len):
            # 使用最后一个时间步的特征进行预测
            pred_input = current_input[:, -1:, :]  # [batch_size, 1, d_model]
            pred_output = self.output_projection(pred_input)  # [batch_size, 1, 1]
            predictions.append(pred_output)
            
            # 更新输入（在实际应用中，这里应该结合最新的特征数据）
            if step < pred_len - 1:
                # 为下一步预测准备输入（简化版本）
                next_features = pred_input  # 实际应用中需要结合新的环境特征
                current_input = torch.cat([current_input[:, 1:, :], next_features], dim=1)
        
        return torch.cat(predictions, dim=1)  # [batch_size, pred_len, 1]


class RealTimePredictor:
    """实时预测器"""
    
    def __init__(self, model: RealTimeHVACTransformer, config: HVACForecastConfig):
        self.model = model
        self.config = config
        self.model.eval()
        
        # 数据缓存
        self.load_buffer = []
        self.numerical_buffer = []
        self.time_buffer = []
        
    def update_data(self, new_load: float, new_numerical: np.ndarray, 
                   timestamp: datetime.datetime):
        """更新实时数据"""
        # 提取时间特征
        time_features = [
            timestamp.hour,
            timestamp.weekday(),
            timestamp.month - 1,
            (timestamp.month - 1) // 3  # 季节
        ]
        
        # 添加到缓冲区
        self.load_buffer.append(new_load)
        self.numerical_buffer.append(new_numerical)
        self.time_buffer.append(time_features)
        
        # 保持缓冲区大小
        if len(self.load_buffer) > self.config.seq_len:
            self.load_buffer.pop(0)
            self.numerical_buffer.pop(0)
            self.time_buffer.pop(0)
    
    def predict_remaining_day(self) -> Tuple[np.ndarray, Dict]:
        """预测今日剩余时间的负荷"""
        if len(self.load_buffer) < self.config.seq_len:
            raise ValueError(f"需要至少 {self.config.seq_len} 个历史数据点")
        
        # 准备输入数据
        load_data = torch.tensor(self.load_buffer[-self.config.seq_len:], 
                                dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
        numerical_features = torch.tensor(self.numerical_buffer[-self.config.seq_len:], 
                                        dtype=torch.float32).unsqueeze(0)
        time_features = torch.tensor(self.time_buffer[-self.config.seq_len:], 
                                   dtype=torch.long).unsqueeze(0)
        
        # 移动到设备
        load_data = load_data.to(self.config.device)
        numerical_features = numerical_features.to(self.config.device)
        time_features = time_features.to(self.config.device)
        
        # 计算今日剩余小时数
        current_hour = datetime.datetime.now().hour
        remaining_hours = 24 - current_hour
        
        # 进行预测
        with torch.no_grad():
            predictions = self.model(load_data, numerical_features, time_features, remaining_hours)
            predictions_np = predictions.cpu().numpy().squeeze()
        
        # 生成时间戳
        current_time = datetime.datetime.now()
        time_stamps = [current_time + datetime.timedelta(hours=i+1) 
                      for i in range(remaining_hours)]
        
        # 计算预测统计信息
        stats = {
            'mean_load': float(np.mean(predictions_np)),
            'max_load': float(np.max(predictions_np)),
            'min_load': float(np.min(predictions_np)),
            'peak_hour': time_stamps[np.argmax(predictions_np)],
            'total_consumption': float(np.sum(predictions_np)),
            'remaining_hours': remaining_hours
        }
        
        return predictions_np, stats
    
    def get_next_hour_prediction(self) -> float:
        """获取下一小时的负荷预测"""
        predictions, _ = self.predict_remaining_day()
        return float(predictions[0]) if len(predictions) > 0 else 0.0


# 示例使用
def create_hvac_model():
    """创建空调负荷预测模型"""
    config = HVACForecastConfig(
        d_model=128,
        n_heads=8,
        n_layers=6,
        seq_len=168,  # 一周数据
        max_pred_len=24,
        n_numerical_features=6,  # 温度、湿度、风速、人员数量、设备状态、电价
        dropout=0.1
    )
    
    model = RealTimeHVACTransformer(config)
    predictor = RealTimePredictor(model, config)
    
    return model, predictor, config


if __name__ == "__main__":
    # 创建模型
    model, predictor, config = create_hvac_model()
    
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
    print(f"模型配置: {config}")
    
    # 模拟实时预测
    print("\n=== 实时空调负荷预测系统 ===")
    print("模型已就绪，可以开始实时预测...")
